"""
Simple plotting for GCR Solver Manager.

Basic matplotlib-based plotting for convergence analysis and performance trends.
"""

import logging
from pathlib import Path
from typing import Dict, List, Optional, Any, Tuple
from datetime import datetime

try:
    import matplotlib
    matplotlib.use('Agg')  # Use non-interactive backend
    import matplotlib.pyplot as plt
    import matplotlib.dates as mdates
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False

from ..database.simple_db import SimpleDatabase
from ..logs.manager import LogManager

logger = logging.getLogger(__name__)


class SimplePlotter:
    """Simple plotting tools for GCR solver data."""
    
    def __init__(self, db: Optional[SimpleDatabase] = None, log_manager: Optional[LogManager] = None):
        """Initialize plotter."""
        if not MATPLOTLIB_AVAILABLE:
            raise ImportError("matplotlib is required for plotting. Install with: pip install matplotlib")
        
        self.db = db or SimpleDatabase()
        self.log_manager = log_manager or LogManager()
    
    def plot_convergence_comparison(self, run_ids: List[int], output_path: str, 
                                  title: Optional[str] = None) -> bool:
        """
        Plot convergence comparison between multiple runs.
        
        Args:
            run_ids: List of run IDs to compare
            output_path: Path to save plot
            title: Optional plot title
            
        Returns:
            True if plot created successfully
        """
        try:
            plt.figure(figsize=(12, 8))
            
            colors = ['blue', 'red', 'green', 'orange', 'purple', 'brown', 'pink', 'gray']
            
            for i, run_id in enumerate(run_ids[:8]):  # Limit to 8 runs for readability
                run_data = self.db.get_by_id(run_id)
                if not run_data or run_data['type'] != 'test':
                    continue
                
                results = run_data.get('results', {})
                norms = results.get('norms', {})
                
                if not norms:
                    continue
                
                # Extract norm values (assume they represent convergence over iterations)
                norm_values = list(norms.values())
                iterations = list(range(1, len(norm_values) + 1))
                
                color = colors[i % len(colors)]
                label = f"Run {run_id} ({run_data['solver']}, {run_data['gpu'] or 'cpu'})"
                
                plt.semilogy(iterations, norm_values, color=color, marker='o', 
                           label=label, linewidth=2, markersize=4)
            
            plt.xlabel('Iteration')
            plt.ylabel('Residual Norm')
            plt.title(title or 'Convergence Comparison')
            plt.legend()
            plt.grid(True, alpha=0.3)
            
            # Save plot
            output_file = Path(output_path)
            output_file.parent.mkdir(parents=True, exist_ok=True)
            plt.savefig(output_file, dpi=300, bbox_inches='tight')
            plt.close()
            
            logger.info(f"Convergence plot saved: {output_path}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to create convergence plot: {e}")
            return False
    
    def plot_performance_trends(self, solver: str, gpu: Optional[str] = None, 
                              days: int = 30, output_path: str = None) -> bool:
        """
        Plot performance trends over time.
        
        Args:
            solver: Solver type to analyze
            gpu: GPU type (optional)
            days: Number of days to analyze
            output_path: Path to save plot (auto-generated if None)
            
        Returns:
            True if plot created successfully
        """
        try:
            # Get recent test runs
            recent_runs = self.db.get_latest(100, 'test', solver, gpu)
            
            if not recent_runs:
                logger.warning(f"No test runs found for {solver} {gpu or 'cpu'}")
                return False
            
            # Extract data for plotting
            dates = []
            final_norms = []
            statuses = []
            
            for run in recent_runs:
                try:
                    # Parse date
                    created_at = datetime.fromisoformat(run['created_at'].replace('Z', '+00:00'))
                    dates.append(created_at)
                    
                    # Extract final norm
                    results = run.get('results', {})
                    norms = results.get('norms', {})
                    
                    if norms:
                        # Use the last norm value as final norm
                        final_norm = list(norms.values())[-1]
                        final_norms.append(final_norm)
                    else:
                        final_norms.append(None)
                    
                    statuses.append(run['status'])
                    
                except Exception as e:
                    logger.warning(f"Failed to process run {run['id']}: {e}")
                    continue
            
            if not dates:
                logger.warning("No valid data points for plotting")
                return False
            
            # Create plot
            plt.figure(figsize=(14, 8))
            
            # Plot successful runs
            success_dates = [d for d, s, n in zip(dates, statuses, final_norms) 
                           if s == 'success' and n is not None]
            success_norms = [n for d, s, n in zip(dates, statuses, final_norms) 
                           if s == 'success' and n is not None]
            
            if success_dates:
                plt.semilogy(success_dates, success_norms, 'go-', label='Successful', 
                           markersize=6, linewidth=2)
            
            # Plot failed runs
            failed_dates = [d for d, s, n in zip(dates, statuses, final_norms) 
                          if s == 'failed' and n is not None]
            failed_norms = [n for d, s, n in zip(dates, statuses, final_norms) 
                          if s == 'failed' and n is not None]
            
            if failed_dates:
                plt.semilogy(failed_dates, failed_norms, 'rx', label='Failed', 
                           markersize=8, markeredgewidth=2)
            
            plt.xlabel('Date')
            plt.ylabel('Final Residual Norm')
            plt.title(f'Performance Trends: {solver} ({gpu or "cpu"}) - Last {days} Days')
            plt.legend()
            plt.grid(True, alpha=0.3)
            
            # Format x-axis
            plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%m/%d'))
            plt.gca().xaxis.set_major_locator(mdates.DayLocator(interval=max(1, days//10)))
            plt.xticks(rotation=45)
            
            # Save plot
            if not output_path:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                output_path = f"performance_trends_{solver}_{gpu or 'cpu'}_{timestamp}.png"
            
            output_file = Path(output_path)
            output_file.parent.mkdir(parents=True, exist_ok=True)
            plt.savefig(output_file, dpi=300, bbox_inches='tight')
            plt.close()
            
            logger.info(f"Performance trends plot saved: {output_path}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to create performance trends plot: {e}")
            return False
    
    def plot_solver_comparison(self, solvers: List[str], gpu: Optional[str] = None, 
                              metric: str = 'final_norm', output_path: str = None) -> bool:
        """
        Plot comparison between different solvers.
        
        Args:
            solvers: List of solver types to compare
            gpu: GPU type (optional)
            metric: Metric to compare ('final_norm', 'convergence_rate')
            output_path: Path to save plot
            
        Returns:
            True if plot created successfully
        """
        try:
            solver_data = {}
            
            # Collect data for each solver
            for solver in solvers:
                recent_runs = self.db.get_latest(20, 'test', solver, gpu)
                successful_runs = [run for run in recent_runs if run['status'] == 'success']
                
                if not successful_runs:
                    continue
                
                values = []
                for run in successful_runs:
                    results = run.get('results', {})
                    norms = results.get('norms', {})
                    
                    if norms:
                        if metric == 'final_norm':
                            value = list(norms.values())[-1]
                        else:  # convergence_rate - simplified as norm improvement
                            norm_values = list(norms.values())
                            if len(norm_values) > 1:
                                value = norm_values[0] / norm_values[-1]  # Improvement ratio
                            else:
                                value = 1.0
                        
                        values.append(value)
                
                if values:
                    solver_data[solver] = values
            
            if not solver_data:
                logger.warning("No data available for solver comparison")
                return False
            
            # Create box plot
            plt.figure(figsize=(12, 8))
            
            data_for_plot = []
            labels = []
            
            for solver, values in solver_data.items():
                data_for_plot.append(values)
                labels.append(f"{solver}\n({len(values)} runs)")
            
            bp = plt.boxplot(data_for_plot, labels=labels, patch_artist=True)
            
            # Color boxes
            colors = ['lightblue', 'lightgreen', 'lightcoral', 'lightyellow', 'lightpink']
            for patch, color in zip(bp['boxes'], colors):
                patch.set_facecolor(color)
            
            if metric == 'final_norm':
                plt.yscale('log')
                plt.ylabel('Final Residual Norm')
                plt.title(f'Solver Comparison: Final Residual Norms ({gpu or "cpu"})')
            else:
                plt.ylabel('Convergence Rate (Initial/Final Norm)')
                plt.title(f'Solver Comparison: Convergence Rates ({gpu or "cpu"})')
            
            plt.xlabel('Solver')
            plt.grid(True, alpha=0.3)
            
            # Save plot
            if not output_path:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                output_path = f"solver_comparison_{metric}_{gpu or 'cpu'}_{timestamp}.png"
            
            output_file = Path(output_path)
            output_file.parent.mkdir(parents=True, exist_ok=True)
            plt.savefig(output_file, dpi=300, bbox_inches='tight')
            plt.close()
            
            logger.info(f"Solver comparison plot saved: {output_path}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to create solver comparison plot: {e}")
            return False
    
    def plot_build_success_rate(self, days: int = 30, output_path: str = None) -> bool:
        """
        Plot build success rate over time.
        
        Args:
            days: Number of days to analyze
            output_path: Path to save plot
            
        Returns:
            True if plot created successfully
        """
        try:
            # Get recent build runs
            recent_builds = self.db.get_latest(200, 'build')
            
            if not recent_builds:
                logger.warning("No build runs found")
                return False
            
            # Group by day and calculate success rates
            daily_stats = {}
            
            for build in recent_builds:
                try:
                    date = build['created_at'][:10]  # YYYY-MM-DD
                    
                    if date not in daily_stats:
                        daily_stats[date] = {'total': 0, 'successful': 0}
                    
                    daily_stats[date]['total'] += 1
                    if build['status'] == 'success':
                        daily_stats[date]['successful'] += 1
                        
                except Exception as e:
                    logger.warning(f"Failed to process build {build['id']}: {e}")
                    continue
            
            if not daily_stats:
                logger.warning("No valid build data for plotting")
                return False
            
            # Sort by date and calculate success rates
            sorted_dates = sorted(daily_stats.keys())
            dates = [datetime.strptime(d, '%Y-%m-%d') for d in sorted_dates]
            success_rates = [daily_stats[d]['successful'] / daily_stats[d]['total'] * 100 
                           for d in sorted_dates]
            total_builds = [daily_stats[d]['total'] for d in sorted_dates]
            
            # Create plot
            fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(14, 10))
            
            # Success rate plot
            ax1.plot(dates, success_rates, 'b-o', linewidth=2, markersize=6)
            ax1.set_ylabel('Success Rate (%)')
            ax1.set_title(f'Build Success Rate Over Time (Last {days} Days)')
            ax1.grid(True, alpha=0.3)
            ax1.set_ylim(0, 100)
            
            # Total builds plot
            ax2.bar(dates, total_builds, alpha=0.7, color='green')
            ax2.set_ylabel('Total Builds per Day')
            ax2.set_xlabel('Date')
            ax2.set_title('Build Activity Over Time')
            ax2.grid(True, alpha=0.3)
            
            # Format x-axis for both subplots
            for ax in [ax1, ax2]:
                ax.xaxis.set_major_formatter(mdates.DateFormatter('%m/%d'))
                ax.xaxis.set_major_locator(mdates.DayLocator(interval=max(1, len(dates)//10)))
                plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
            
            plt.tight_layout()
            
            # Save plot
            if not output_path:
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                output_path = f"build_success_rate_{timestamp}.png"
            
            output_file = Path(output_path)
            output_file.parent.mkdir(parents=True, exist_ok=True)
            plt.savefig(output_file, dpi=300, bbox_inches='tight')
            plt.close()
            
            logger.info(f"Build success rate plot saved: {output_path}")
            return True
            
        except Exception as e:
            logger.error(f"Failed to create build success rate plot: {e}")
            return False